/**@file
   WDF driver helper class.
   
   @par license 
   Copyright 2012 Lu, Ken. All rights reserved.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 
   @author: Lu, Ken (bluewish.ken.lu@gmail.com)
**/

#include <stdexcept>
#include <fstream>
#include <assert.h>

#include "logger.h"
#include "wdfhelper.h"

namespace ken {

uint16_t WdfDriverHelper::version_major_ = 1;
uint16_t WdfDriverHelper::version_minor_ = 9;

WdfDriverHelper::WdfDriverHelper(void)
{
    hCoInstallerDll_            = NULL;
    pfn_WdfPreDeviceInstall_    = NULL;
    pfn_WdfPostDeviceInstall_   = NULL;
    pfn_WdfPreDeviceRemove_     = NULL;
    pfn_WdfPostDeviceRemove_    = NULL;
}

WdfDriverHelper::~WdfDriverHelper(void)
{
    if (hCoInstallerDll_ != NULL) {
        FreeLibrary(hCoInstallerDll_);
    }
}

WdfDriverHelper::WdfDriverHelper(
    const wchar_t *device_name, 
    const wchar_t* driver_name
    )
{
    WdfDriverHelper();
    device_name_ = device_name;
    driver_name_ = driver_name;

}

void WdfDriverHelper::SetVersion(uint16_t major, uint16_t minor)
{
    WdfDriverHelper::version_major_ = major;
    WdfDriverHelper::version_minor_ = minor;
}

void WdfDriverHelper::Initialize() 
{
    wchar_t  dll_path_buffer[MAX_PATH] = {0};
    std::wstring dll_path_str;

    ::GetCurrentDirectoryW(MAX_PATH, dll_path_buffer);
    dll_path_str = std::wstring(dll_path_buffer);
    
    dll_path_str += L"\\";
    dll_path_str += getCoInstallerDllName();

    hCoInstallerDll_ = ::LoadLibraryW(dll_path_str.c_str());
    pfn_WdfPreDeviceInstall_ = 
        (PFN_WDFPREDEVICEINSTALL)::GetProcAddress(
                                    hCoInstallerDll_, 
                                     "WdfPreDeviceInstall");
    pfn_WdfPostDeviceInstall_ = 
        (PFN_WDFPOSTDEVICEINSTALL)::GetProcAddress(
                                    hCoInstallerDll_, 
                                    "WdfPostDeviceInstall");
    pfn_WdfPreDeviceRemove_ = 
        (PFN_WDFPREDEVICEREMOVE)::GetProcAddress(
                                    hCoInstallerDll_, 
                                    "WdfPreDeviceRemove");
    pfn_WdfPostDeviceRemove_ = 
        (PFN_WDFPOSTDEVICEREMOVE)::GetProcAddress(
                                    hCoInstallerDll_, 
                                    "WdfPostDeviceRemove");
    if (pfn_WdfPreDeviceInstall_ == NULL    ||
        pfn_WdfPostDeviceInstall_ == NULL   ||
        pfn_WdfPreDeviceRemove_ == NULL     ||
        pfn_WdfPostDeviceRemove_ == NULL ) {
        FreeLibrary(hCoInstallerDll_);
        hCoInstallerDll_ = NULL;
        throw std::runtime_error("Fail to get function from WdfCoInstaller01009.dll");
    }
}

void WdfDriverHelper::Install(
    const wchar_t* driver_dir, 
    const wchar_t* inf_name, 
    const wchar_t* inf_wdf_section_name
    )
{
    std::wstring inf_path       = getInfFilePath(driver_dir, inf_name);
    std::wstring driver_path    = getDriverFilePath(driver_dir);

    //
    // pre install
    //
    assert(pfn_WdfPreDeviceInstall_ != NULL);
    uint32_t err = pfn_WdfPreDeviceInstall_(inf_path.c_str(), inf_wdf_section_name);
    if (err != ERROR_SUCCESS) {
        if (err == ERROR_SUCCESS_REBOOT_REQUIRED) {
            throw std::runtime_error("System nees reboot before the driver installation can be processed.");
        } else {
            char err_msg[256] = {0};
            sprintf_s(err_msg, "Fail to process INF's wdf section before installation. error=0x%x", err);
            throw std::runtime_error(err_msg);
        }
    }

    //
    // install
    //
    SC_HANDLE sc_handle = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (sc_handle == NULL) {
        ::CloseServiceHandle(sc_handle);
        reportWindowsError("Fail to open sc manager.");
    }

    SC_HANDLE service_handle = ::CreateServiceW(
        sc_handle, 
        driver_name_.c_str(), 
        driver_name_.c_str(), 
        SERVICE_ALL_ACCESS, 
        SERVICE_KERNEL_DRIVER, 
        SERVICE_DEMAND_START, 
        SERVICE_ERROR_NORMAL, 
        driver_path.c_str(), 
        NULL, NULL, NULL, NULL, NULL);
    if (service_handle == NULL) {
        ::CloseServiceHandle(sc_handle);
        if (::GetLastError() == ERROR_SERVICE_EXISTS) {
            return;
        }
        reportWindowsError("Fail to create service.");
    }
    ::CloseServiceHandle(service_handle);
    ::CloseServiceHandle(sc_handle);


    //
    // post install
    //
    assert(pfn_WdfPostDeviceInstall_ != NULL);
    err = pfn_WdfPostDeviceInstall_(inf_path.c_str(), inf_wdf_section_name);
    if (err != ERROR_SUCCESS) {
        char err_msg[256] = {0};
        sprintf_s(err_msg, "Fail to process INF's wdf section after installation. error=0x%x", err);
        throw std::runtime_error(err_msg);
    }
}

void WdfDriverHelper::Start()
{
    SC_HANDLE sc_handle = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (sc_handle == NULL) {
        reportWindowsError("Fail to open sc manager.");
    }

    SC_HANDLE service_handle = ::OpenService(
        sc_handle,
        driver_name_.c_str(),
        SERVICE_ALL_ACCESS
        );
    if (service_handle == NULL) {
        ::CloseServiceHandle(sc_handle);
        reportWindowsError("Fail to open service.");
    }

    bool ret = ::StartServiceA(service_handle, 0, NULL) == TRUE;
    ::CloseServiceHandle(service_handle);
    ::CloseServiceHandle(sc_handle);

    if (!ret) {
        if (::GetLastError() != ERROR_SERVICE_ALREADY_RUNNING) {
            reportWindowsError("Fail to start service.");
        }
    }
}

void WdfDriverHelper::Remove(
    const wchar_t* driver_dir, 
    const wchar_t* inf_name, 
    const wchar_t* inf_wdf_section_name
    )
{
    std::wstring inf_path       = getInfFilePath(driver_dir, inf_name);

    //
    // pre remove
    //
    assert(pfn_WdfPreDeviceRemove_ != NULL);
    uint32_t err = pfn_WdfPreDeviceRemove_(inf_path.c_str(), inf_wdf_section_name);
    if (err != ERROR_SUCCESS) {
        char err_msg[256] = {0};
        sprintf_s(err_msg, "Fail to process INF's wdf section before remove. error=0x%x", err);
        throw std::runtime_error(err_msg);
    }

    //
    // install
    //
    SC_HANDLE sc_handle = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (sc_handle == NULL) {
        reportWindowsError("Fail to open sc manager.");
    }

    SC_HANDLE service_handle = ::OpenService(
        sc_handle,
        driver_name_.c_str(),
        SERVICE_ALL_ACCESS
        );
    if (service_handle == NULL) {
        ::CloseServiceHandle(sc_handle);
        reportWindowsError("Fail to open service.");
    }

    bool ret = ::DeleteService(service_handle) == TRUE;
    ::CloseServiceHandle(service_handle);
    ::CloseServiceHandle(sc_handle);

    //
    // post remove
    //
    assert(pfn_WdfPostDeviceRemove_ != NULL);
    err = pfn_WdfPostDeviceRemove_(inf_path.c_str(), inf_wdf_section_name);
    if (err != ERROR_SUCCESS) {
        char err_msg[256] = {0};
        sprintf_s(err_msg, "Fail to process INF's wdf section after remove. error=0x%x", err);
        throw std::runtime_error(err_msg);
    }

    if (!ret) {
        throw std::runtime_error("Fail to remove service!");
    }
}

void WdfDriverHelper::Stop()
{
    SC_HANDLE sc_handle = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (sc_handle == NULL) {
        reportWindowsError("Fail to open sc manager");
    }

    SC_HANDLE service_handle = ::OpenService(
        sc_handle,
        driver_name_.c_str(),
        SERVICE_ALL_ACCESS
        );
    if (service_handle == NULL) {
        ::CloseServiceHandle(sc_handle);
        reportWindowsError("Fail to open service");
    }

    SERVICE_STATUS service_status;
    bool ret = ::ControlService(service_handle, SERVICE_CONTROL_STOP, &service_status) == TRUE;

    if (!ret) {
        char mesg[256] = {0};
        sprintf_s(mesg, "Fail to stop service, status=0x%x", service_status);
        reportWindowsError(mesg);
    }

    ::CloseServiceHandle(service_handle);
    ::CloseServiceHandle(sc_handle);
}

std::wstring WdfDriverHelper::getCoInstallerDllName()
{
    wchar_t name[MAX_PATH] = {0};
    swprintf_s(name, L"WdfCoInstaller%02d%03d.dll", WdfDriverHelper::version_major_, WdfDriverHelper::version_minor_);
    return name;
}

bool WdfDriverHelper::checkFileExist(std::wstring& file_path)
{
    fstream check_file;
    bool ret = true;

    check_file.open(file_path, ios::in);
    if (!check_file) {
        ret = false;
    }
    check_file.close();

    return ret;
}

void WdfDriverHelper::reportWindowsError(const char* err)
{
    char mesg[256] = {0};
    sprintf_s(mesg, "%s, error code=0x%x", err, ::GetLastError());
    throw std::runtime_error(mesg);
}

std::wstring WdfDriverHelper::getDriverFilePath(const wchar_t *driver_dir)
{
    std::wstring path_str;

    if (driver_dir == NULL) {
        wchar_t path[MAX_PATH] = {0};
        ::GetCurrentDirectoryW(MAX_PATH, path);
        path_str = path;
    } else {
        path_str = driver_dir;
    }

    std::wstring driver_full_path = path_str + L"\\" + driver_name_ + L".sys";
    if (!checkFileExist(driver_full_path)) {
        throw std::runtime_error("Driver file does not exist!");
    }

    return driver_full_path;
}

std::wstring WdfDriverHelper::getInfFilePath(
    const wchar_t* driver_dir, 
    const wchar_t* inf_name
    )
{
    std::wstring path_str;

    if (driver_dir == NULL) {
        wchar_t path[MAX_PATH] = {0};
        ::GetCurrentDirectoryW(MAX_PATH, path);
        path_str = path;
    } else {
        path_str = driver_dir;
    }

    std::wstring inf_full_path = path_str + L"\\" + inf_name;
    if (!checkFileExist(inf_full_path)) {
        throw std::runtime_error("Inf file does not exist!");
    }

    return inf_full_path;
}

uint8_t* WdfDriverHelper::u2a(const uint16_t *src, uint8_t *target)
{
    assert(src != NULL);
    assert(target != NULL);

    uint16_t *p = (uint16_t*)src;
    uint8_t *q = target;
    while (*p != 0) {
        *(q++) = (uint8_t)*(p++);
    }
    *q = 0;
    return target;
}

}